package com.util.orm.dao;

import org.apache.ibatis.session.SqlSessionFactory;
import org.mybatis.spring.SqlSessionFactoryBean;
import org.mybatis.spring.SqlSessionTemplate;
import org.springframework.core.io.Resource;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

public class SqlSessionTemplateRouter {

    private Map<String, SqlSessionTemplate> sqlSessionTemplateMap;

    public SqlSessionTemplateRouter(Map<String, DataSource> targetDataSources,
                                    DataSource defaultTargetDataSource,
                                    Resource[] mapperLocations) {
        this.sqlSessionTemplateMap = new HashMap<>();
        synchronized (SqlSessionTemplateRouter.class) {
            targetDataSources.put(null, defaultTargetDataSource); //默认数据源用 null 作为 key
            for (Map.Entry<String, DataSource> entry : targetDataSources.entrySet()) {
                SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
                sqlSessionFactoryBean.setDataSource(entry.getValue());
                sqlSessionFactoryBean.setMapperLocations(mapperLocations);
                try {
                    SqlSessionFactory sqlSessionFactory = sqlSessionFactoryBean.getObject();
                    sqlSessionTemplateMap.put(entry.getKey(), new SqlSessionTemplate(sqlSessionFactory));
                } catch (Exception e) {
                    throw new RuntimeException("Get sqlSessionFactory fail! " + e);
                }
            }
        }
    }

    /**
     * 根据分库 key 查找对应 SqlSessionTemplate
     * 若 key 为 null，则返回默认 SqlSessionTemplate
     * @param key
     * @return
     */
    public SqlSessionTemplate getSqlSessionTemplate(String key) {
        SqlSessionTemplate sqlSessionTemplate = sqlSessionTemplateMap.get(key);
        if (null == sqlSessionTemplate) {
            throw new RuntimeException(String.format("SqlSessionTemplate of key %s does not exist", key));
        }
        return sqlSessionTemplate;
    }
}
